package uk.ac.strath.cis.metric_spaces.indices;

import java.util.ArrayList;
import java.util.Collections;
import java.util.Comparator;
import java.util.List;
import java.util.Set;

import uk.ac.strath.cis.metric_spaces.framework.AbstractMetricIndex;
import uk.ac.strath.cis.metric_spaces.framework.NearestNeighbourQuery;

/**
 * @author Robert
 *
 * @param <E>
 */
public class DiscreteNearestNeighbourQuery<E> implements
        NearestNeighbourQuery<E> {

    private AbstractMetricIndex<E> index;

    /* (non-Javadoc)
     * @see uk.ac.strath.cis.rmoss.metric_spaces.framework.NearestNeighbourQuery#setIndex(uk.ac.strath.cis.rmoss.metric_spaces.framework.AbstractMetricIndex)
     */
    public void setIndex(AbstractMetricIndex<E> index) {
        this.index = index;
    }

    /* (non-Javadoc)
     * @see uk.ac.strath.cis.rmoss.metric_spaces.framework.NearestNeighbourQuery#query(java.lang.Object, int)
     */
    public Set<E> query(final E element, int number) {
        Set<E> result;
        int i = 1;
        while ((result = index.rangeQuery(element, i)).size() < number
                && index.getSize() >= i) {
            i++;
        }
        
        number = (number > result.size()) ? result.size() : number;
        List<E> temp = new ArrayList<E>();
        temp.addAll(result);
        Collections.sort(temp, new Comparator<E>() {

            public int compare(E o1, E o2) {
                double d1 = index.distance(element, o1);
                double d2 = index.distance(element, o2);
                return (d1 > d2) ? 1 : (d1 == d2) ? 0 : -1;
            }
        });
        result.clear();
        result.addAll(temp.subList(0, number));
        return result;
    }

}
